# -*- coding: utf-8 -*-
"""
Author: Qwen
Date: 2025/09/18
File: app_to_server_wake_up_v1.py
功能：基于 asyncio 的异步 UDP 压力测试，低内存、高并发，支持 QPS 限流与每分钟统计
"""

import struct
import asyncio
import random
import time
from datetime import datetime

# ==================== 配置参数 ====================
SERVER_IP = "47.107.56.10"
SERVER_PORT = 50001
UDP_TIMEOUT = 5  # 秒
TOTAL_DURATION = 24 * 3600  # 24 小时

# QPS 控制
TARGET_QPS = 5000  # 每秒最多发送多少请求
INTERVAL_PER_REQUEST = 1.0 / TARGET_QPS  # 平均每请求间隔（秒）

# DID 范围
DID_PREFIX = "IOTGDD-"
DID_SUFFIX = "-EDIDSC"
DID_MIN_INDEX = 0
DID_MAX_INDEX = 60000

# 统计变量（主线程访问，asyncio 中线程安全）
minute_requests = 0
minute_successes = 0
total_requests = 0
total_successes = 0
stats_lock = asyncio.Lock()  # 异步锁

# UDP 协议头
START_CODE = 0xFFFFFFA1
MSG_TYPE = 1
MSG_HEADER_FORMAT = ">III"  # big-endian: StartCode, MsgType, MsgLen


# ==================== 工具函数 ====================
def create_request(did: str) -> bytes:
	"""构建 UDP 请求包"""
	payload = did.encode('utf-8')
	msg_len = len(payload)
	header = struct.pack(MSG_HEADER_FORMAT, START_CODE, MSG_TYPE, msg_len)
	return header + payload


async def print_minute_stats():
	"""每分钟打印一次统计信息"""
	while True:
		await asyncio.sleep(60 - time.time() % 60)  # 对齐整分钟
		
		async with stats_lock:
			reqs = minute_requests
			succ = minute_successes
			minute_requests = 0
			minute_successes = 0
		
		success_rate = (succ / reqs * 100) if reqs > 0 else 0
		qps = reqs / 60.0
		now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
		print(f"[{now}] 分钟统计 | 发送: {reqs:,} | 成功: {succ:,} | "
		      f"成功率: {success_rate:.1f}% | QPS: {qps:.1f}")


async def wait_for_rate_limit():
	"""异步限流：确保不超过 TARGET_QPS"""
	global last_send_time
	now = time.time()
	min_interval = INTERVAL_PER_REQUEST
	
	async with stats_lock:
		# 使用全局变量记录上次发送时间
		if hasattr(wait_for_rate_limit, 'last_time'):
			elapsed = now - wait_for_rate_limit.last_time
			if elapsed < min_interval:
				await asyncio.sleep(min_interval - elapsed)
		wait_for_rate_limit.last_time = time.time()


# ==================== 全局变量 ====================
transport = None  # 复用的 UDP transport
target_addr = (SERVER_IP, SERVER_PORT)


# ==================== 优化后的发送函数 ====================
async def send_single_request():
	"""发送单个 UDP 请求（复用 transport）"""
	global total_requests, total_successes, minute_requests, minute_successes
	
	# 限流：基于时间戳控制发送频率
	now = time.time()
	interval = INTERVAL_PER_REQUEST
	if hasattr(send_single_request, 'last_time'):
		delay = interval - (now - send_single_request.last_time)
		if delay > 0:
			await asyncio.sleep(delay)
	send_single_request.last_time = time.time()
	
	# 随机生成 DID
	idx = random.randint(DID_MIN_INDEX, DID_MAX_INDEX)
	did = f"{DID_PREFIX}{idx:06d}{DID_SUFFIX}"
	request_data = create_request(did)
	
	success = False
	try:
		# 直接通过复用的 transport 发送（无连接创建开销）
		transport.sendto(request_data, target_addr)
		success = True
	except Exception:
		pass  # UDP 不保证可靠，失败忽略
	
	# 更新统计
	async with stats_lock:
		total_requests += 1
		minute_requests += 1
		if success:
			total_successes += 1
			minute_successes += 1


# ==================== 主函数（优化版） ====================
async def main():
	global loop, transport
	loop = asyncio.get_running_loop()
	
	# 创建一次 UDP transport 并复用
	transport, _ = await loop.create_datagram_endpoint(
		lambda: None,
		remote_addr=target_addr
	)
	print(f"UDP transport 已创建并复用: {target_addr}")
	
	print(f"【异步 UDP 压力测试启动】")
	print(f"目标服务器: {SERVER_IP}:{SERVER_PORT}")
	print(f"目标 QPS: {TARGET_QPS:,}")
	print(f"DID 范围: {DID_MIN_INDEX} ~ {DID_MAX_INDEX}")
	print(f"总时长: {TOTAL_DURATION // 3600} 小时")
	print(f"开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
	print("-" * 60)
	
	# 启动统计任务
	asyncio.create_task(print_minute_stats())
	
	# === 核心压测循环：持续发送，不限批次 ===
	start_time = time.time()
	try:
		while time.time() - start_time < TOTAL_DURATION:
			# 不用 gather，直接 fire-and-forget
			asyncio.create_task(send_single_request())
			# 可加轻微 delay 避免事件循环过载
			# await asyncio.sleep(0)  # 让出控制权
	except KeyboardInterrupt:
		pass
	finally:
		transport.close()
	
	# 打印最终统计（同原逻辑）
	run_time = time.time() - start_time
	overall_qps = total_requests / run_time if run_time > 0 else 0
	success_rate = (total_successes / total_requests * 100) if total_requests > 0 else 0
	
	print("=" * 60)
	print("【压力测试完成】")
	print(f"结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
	print(f"总运行时间: {run_time // 3600:.0f}h {(run_time % 3600) // 60:.0f}m")
	print(f"总请求数: {total_requests:,}")
	print(f"成功数: {total_successes:,}")
	print(f"总体成功率: {success_rate:.2f}%")
	print(f"平均 QPS: {overall_qps:.2f}")
	print("=" * 60)


if __name__ == "__main__":
	asyncio.run(main())
